package com.yifeng.repo.controller.rocketmq.message;

import com.gomcarter.frameworks.base.common.AssertUtils;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.yifeng.repo.base.utils.common.BaseUtil;
import com.yifeng.repo.controller.rocketmq.configure.RocketmqProperties;
import com.yifeng.repo.controller.rocketmq.message.client.DefaultMQPushConsumerExt;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.MessageListener;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.consumer.listener.MessageListenerOrderly;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueAveragely;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
import org.apache.rocketmq.common.message.MessageQueue;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

/**
 * Created by daibing on 2023/8/7.
 */
@Slf4j
public class PushConsumeManager {
    protected final ConcurrentMap<String, DefaultMQPushConsumerExt> topic2PushConsumer = new ConcurrentHashMap<>();
    protected final String applicationName;
    protected final RocketmqProperties properties;

    public PushConsumeManager(String applicationName, RocketmqProperties properties) {
        this.applicationName = applicationName;
        this.properties = properties;
    }

    public void init() {
        log.info("yfcloud controller PushConsumeManager init success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    public void destroy() {
        topic2PushConsumer.forEach((k, v) -> {
            v.unsubscribe(k);
            v.shutdown();
        });
        log.info("yfcloud controller PushConsumeManager destroy success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    /**
     * 增加Topic消费
     *
     * @param topic           消息主题
     * @param subExpression   子表达式，比如：全部是 *，多个tag是 TagA || TagC || TagD
     * @param messageListener 消息监听器
     */
    public void addTopicConsumer(String topic, String subExpression, MessageListener messageListener) {
        if (!(messageListener instanceof MessageListenerConcurrently) && !(messageListener instanceof MessageListenerOrderly)) {
            throw new CustomException("当前仅支持顺序监听器和并发监听器：" + messageListener.toString());
        }
        String key = topic + "-" + subExpression;
        DefaultMQPushConsumerExt pushConsumer = topic2PushConsumer.get(key);
        if (pushConsumer != null) {
            return;
        }
        synchronized (topic2PushConsumer) {
            pushConsumer = topic2PushConsumer.get(key);
            if (pushConsumer != null) {
                return;
            }
            pushConsumer = buildDefaultMQPushConsumer(topic, subExpression, messageListener);
            topic2PushConsumer.put(key, pushConsumer);
        }
    }

    /**
     * 移除Topic消费
     */
    public void removeTopicConsumer(String topic, String subExpression) {
        String key = topic + "-" + subExpression;
        DefaultMQPushConsumer pushConsumer = topic2PushConsumer.remove(key);
        if (pushConsumer != null) {
            log.info("removeTopicConsumer success: {}-{}", topic, subExpression);
            pushConsumer.unsubscribe(topic);
            pushConsumer.shutdown();
        }
    }

    public List<String> listConsumerId(String topic, String subExpression) {
        String key = topic + "-" + subExpression;
        DefaultMQPushConsumerExt pushConsumer = topic2PushConsumer.get(key);
        if (pushConsumer == null) {
            return Collections.emptyList();
        }
        String consumerGroup = pushConsumer.getConsumerGroup();
        List<String> consumerIds = pushConsumer.findConsumerIdList(topic, consumerGroup);
        return consumerIds != null ? consumerIds : Collections.emptyList();
    }

    /**
     * 列出当前消息主题的消费subExpression
     */
    public List<String> listConsumerSubExpression(String topic) {
        return topic2PushConsumer.keySet().stream()
                .filter(k -> k.startsWith(topic + "-"))
                .map(s -> s.substring((topic + "-").length()))
                .collect(Collectors.toList());
    }

    /**
     * 查询指定topic的消息队列
     */
    public List<MessageQueue> listMessageQueue(String topic) {
        AssertUtils.notEmpty(topic2PushConsumer.values(), "当前还没有初始化消费者，不能获取消息队列列表！");
        DefaultMQPushConsumer pushConsumer = topic2PushConsumer.entrySet()
                .stream()
                .filter(entry -> entry.getKey().startsWith(topic + "-"))
                .findFirst()
                .orElseThrow(() -> new CustomException("获取不到Topic对应的消费者：" + topic)).getValue();
        try {
            return new ArrayList<>(pushConsumer.fetchSubscribeMessageQueues(topic));
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
    }

    protected DefaultMQPushConsumerExt buildDefaultMQPushConsumer(String topic, String subExpression, MessageListener messageListener) {
        // 获取topic配置
        RocketmqProperties.TopicProperties topicProperties = getTopicProperties(properties, topic);
        String consumerGroup = getConsumerGroup(topicProperties, subExpression);

        // build push consumer
        AclClientRPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(topicProperties.getAccessKey(), topicProperties.getSecretKey()));
        DefaultMQPushConsumerExt pushConsumer = new DefaultMQPushConsumerExt(consumerGroup, rpcHook, new AllocateMessageQueueAveragely());
        pushConsumer.setInstanceName(applicationName + "-" + topic + "-push-subscribe-" + subExpression + "-" + consumerGroup);
        pushConsumer.setNamesrvAddr(topicProperties.getNamesrvAddr());
        pushConsumer.setMaxReconsumeTimes(topicProperties.getMaxReconsumeTimes());
        pushConsumer.setConsumeFromWhere(topicProperties.getConsumeFromWhere());
        pushConsumer.setConsumeThreadMin(topicProperties.getPushConsumeThreadMin());
        pushConsumer.setConsumeThreadMax(topicProperties.getPushConsumeThreadMax());
        pushConsumer.setPullBatchSize(topicProperties.getPushBatchSize());
        pushConsumer.setConsumeMessageBatchMaxSize(topicProperties.getConsumeMessageBatchMaxSize());
        pushConsumer.setSuspendCurrentQueueTimeMillis(3000);
        if (messageListener instanceof MessageListenerConcurrently) {
            pushConsumer.registerMessageListener((MessageListenerConcurrently) messageListener);
        } else {
            pushConsumer.registerMessageListener((MessageListenerOrderly) messageListener);
        }

        // consumer start.
        String key = topic + "-" + subExpression;
        try {
            pushConsumer.subscribe(topic, subExpression);
            pushConsumer.start();
            log.info("buildDefaultMQPushConsumer success: {}, consumerGroup={}, consumeThreadMin={}, consumeThreadMax={}, pullBatchSize={}, consumeMessageBatchMaxSize={}",
                    key, pushConsumer.getConsumerGroup(), pushConsumer.getConsumeThreadMin(), pushConsumer.getConsumeThreadMax(),
                    pushConsumer.getPullBatchSize(), pushConsumer.getConsumeMessageBatchMaxSize());
            return pushConsumer;
        } catch (MQClientException e) {
            throw new CustomException("buildDefaultMQPushConsumer failed " + key + "-" + consumerGroup, e);
        }
    }

    protected RocketmqProperties.TopicProperties getTopicProperties(RocketmqProperties properties, String topic) {
        return properties.getTopicPropertiesList()
                .stream()
                .filter(s -> s.getTopic().equals(topic))
                .findFirst()
                .orElseThrow(() -> new CustomException("没有Topic配置：" + topic));
    }

    /**
     * 指定订阅表达式使用指定消费组，如果没有指定订阅表达式，或者如果没有配置，都使用默认消费组
     */
    protected String getConsumerGroup(RocketmqProperties.TopicProperties topicProperties, String subExpression) {
        return Strings.isNullOrEmpty(subExpression) || topicProperties.getConsumerGroupMapping() == null
                ? topicProperties.getConsumerGroup()
                : topicProperties.getConsumerGroupMapping().getOrDefault(subExpression, topicProperties.getConsumerGroup());
    }
}
